import akshare as ak
import pandas as pd
import os
import time
import logging
from datetime import datetime
import warnings

# 忽略警告信息
warnings.filterwarnings('ignore')

# 配置日志系统
def setup_logging():
    # 创建日志目录
    log_dir = 'logs'
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    # 生成日志文件名
    log_file = os.path.join(log_dir, f'stock_data_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log')
    
    # 配置根日志记录器
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s:%(lineno)d - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S',
        handlers=[
            logging.FileHandler(log_file, encoding='utf-8'),
            logging.StreamHandler()
        ]
    )
    
    return logging.getLogger(__name__)

# 初始化日志记录器
logger = setup_logging()

# 处理市场数据列名的辅助函数
def process_market_data_columns(market_df):
    """
    处理市场数据的列名，确保包含标准列名
    """
    # 记录原始列名
    logger.debug(f"原始市场数据列: {list(market_df.columns)}")
    
    # 定义列名映射关系
    column_mapping = {
        '代码': ['代码', 'code', '股票代码', '证券代码'],
        '现价': ['现价', 'price', '最新价', 'current_price'],
        '涨跌幅': ['涨跌幅', 'change_pct', '涨跌幅%', 'percent'],
        '成交量': ['成交量', 'volume', 'volume_traded'],
        '成交额': ['成交额', 'amount', 'turnover_amount'],
        '总市值': ['总市值', 'total_market_cap', '总市值(亿)'],
        '流通市值': ['流通市值', 'circulation_market_cap', '流通市值(亿)'],
        '换手率': ['换手率', 'turnover_rate', '换手率%']
    }
    
    # 重命名列
    for standard_col, possible_cols in column_mapping.items():
        for col in possible_cols:
            if col in market_df.columns:
                if col != standard_col:
                    logger.debug(f"将列 '{col}' 重命名为 '{standard_col}'")
                market_df = market_df.rename(columns={col: standard_col})
                break
    
    logger.debug(f"处理后的市场数据列: {list(market_df.columns)}")
    return market_df

# 获取个股的市值和股价数据
def get_stock_quote_and_market_cap(stock_code):
    """
    获取单个股票的市值和股价数据
    """
    try:
        logger.debug(f"获取股票 {stock_code} 的实时行情数据")
        # 获取实时行情数据
        quote_data = ak.stock_zh_a_spot_em()
        # 处理列名
        quote_data = process_market_data_columns(quote_data)
        # 检查是否有代码列
        if '代码' not in quote_data.columns:
            logger.error("市场数据中未找到代码列")
            return None
        stock_quote = quote_data[quote_data['代码'] == stock_code]
        
        if not stock_quote.empty:
            # 提取所需字段
            result = {'代码': stock_code}
            
            # 要提取的字段列表
            fields = ['现价', '涨跌幅', '成交量', '成交额', '总市值', '流通市值', '换手率']
            
            # 安全地提取每个字段
            for field in fields:
                if field in stock_quote.columns:
                    try:
                        result[field] = stock_quote[field].values[0]
                    except (IndexError, ValueError) as e:
                        logger.warning(f"提取股票 {stock_code} 的 {field} 字段时出错: {e}")
                        result[field] = None
                else:
                    logger.debug(f"股票 {stock_code} 的 {field} 字段不存在")
                    result[field] = None
            logger.debug(f"成功获取股票 {stock_code} 的数据")
            return result
        logger.warning(f"未找到股票 {stock_code} 的数据")
        return None
    except Exception as e:
        logger.error(f"获取股票 {stock_code} 数据时出错: {e}")
        return None

# 批量获取股票数据
def get_batch_stock_data(stock_codes, batch_size=100):
    """
    批量获取股票数据，避免请求过于频繁
    """
    results = []
    total_batches = (len(stock_codes) + batch_size - 1) // batch_size
    
    logger.info(f"开始获取 {len(stock_codes)} 只股票的市值和股价数据...")
    logger.info(f"共分为 {total_batches} 个批次，每批 {batch_size} 只股票")
    
    try:
        # 使用akshare的批量获取方法，更高效
        logger.info("正在获取市场整体行情数据...")
        market_data = ak.stock_zh_a_spot_em()
        logger.info(f"成功获取市场数据，共 {len(market_data)} 条记录")
        
        # 处理市场数据列名
        market_data = process_market_data_columns(market_data)
        
        # 检查是否有代码列
        if '代码' not in market_data.columns:
            logger.error("市场数据中未找到代码列，无法进行批量获取")
            # 直接进入逐条获取的回退模式
            raise KeyError("市场数据中未找到代码列")
        
        # 过滤出我们需要的股票数据
        for i, stock_code in enumerate(stock_codes):
            # 在市场数据中查找当前股票
            stock_data = market_data[market_data['代码'] == stock_code]
            
            if not stock_data.empty:
                # 提取所需字段
                result = {'代码': stock_code}
                
                # 要提取的字段列表
                fields = ['现价', '涨跌幅', '成交量', '成交额', '总市值', '流通市值', '换手率']
                
                # 安全地提取每个字段
                for field in fields:
                    if field in stock_data.columns:
                        try:
                            result[field] = stock_data[field].values[0]
                        except (IndexError, ValueError) as e:
                            logger.warning(f"批量提取股票 {stock_code} 的 {field} 字段时出错: {e}")
                            result[field] = None
                    else:
                        logger.debug(f"批量处理时股票 {stock_code} 的 {field} 字段不存在")
                        result[field] = None
                results.append(result)
            
            # 显示进度
            if (i + 1) % 50 == 0 or (i + 1) == len(stock_codes):
                logger.info(f"已处理 {i + 1}/{len(stock_codes)} 只股票")
        
        logger.info(f"批量获取完成，成功获取 {len(results)} 只股票的数据")
        return pd.DataFrame(results)
        
    except Exception as e:
        logger.error(f"批量获取股票数据时出错: {e}")
        # 回退到逐条获取方式（但效率较低）
        logger.warning("尝试使用逐条获取方式...")
        
        for i, stock_code in enumerate(stock_codes):
            result = get_stock_quote_and_market_cap(stock_code)
            if result:
                results.append(result)
            
            # 显示进度
            if (i + 1) % 50 == 0 or (i + 1) == len(stock_codes):
                logger.info(f"已处理 {i + 1}/{len(stock_codes)} 只股票")
            
            # 避免请求过于频繁
            if (i + 1) % 10 == 0:
                logger.debug("暂停1秒，避免请求过于频繁...")
                time.sleep(1)
        
        logger.info(f"逐条获取完成，成功获取 {len(results)} 只股票的数据")
        return pd.DataFrame(results)

# 处理股票列名的辅助函数
def process_stock_columns(stock_df):
    """
    动态检测并处理股票数据的代码列和名称列
    """
    # 记录原始列名
    logger.debug(f"原始股票数据列: {list(stock_df.columns)}")
    
    # 可能的代码列名列表
    code_column_candidates = ['代码', 'code', '股票代码', '证券代码']
    # 可能的名称列名列表
    name_column_candidates = ['名称', 'name', '股票名称', '证券名称']
    
    # 查找代码列
    code_column_found = None
    for col in code_column_candidates:
        if col in stock_df.columns:
            code_column_found = col
            break
    
    # 如果没找到，尝试使用第一列
    if code_column_found is None and len(stock_df.columns) > 0:
        code_column_found = stock_df.columns[0]
        logger.warning(f"未找到标准代码列，使用第一列 '{code_column_found}' 作为代码列")
    
    # 查找名称列
    name_column_found = None
    for col in name_column_candidates:
        if col in stock_df.columns:
            name_column_found = col
            break
    
    # 如果没找到，尝试使用第二列
    if name_column_found is None and len(stock_df.columns) > 1:
        name_column_found = stock_df.columns[1]
        logger.warning(f"未找到标准名称列，使用第二列 '{name_column_found}' 作为名称列")
    
    # 重命名列
    if code_column_found and code_column_found != '代码':
        logger.debug(f"将列 '{code_column_found}' 重命名为 '代码'")
        stock_df = stock_df.rename(columns={code_column_found: '代码'})
    
    if name_column_found and name_column_found != '名称':
        logger.debug(f"将列 '{name_column_found}' 重命名为 '名称'")
        stock_df = stock_df.rename(columns={name_column_found: '名称'})
    
    # 验证重命名是否成功
    if '代码' not in stock_df.columns or '名称' not in stock_df.columns:
        missing_cols = []
        if '代码' not in stock_df.columns:
            missing_cols.append('代码')
        if '名称' not in stock_df.columns:
            missing_cols.append('名称')
        error_msg = f"无法识别必要的列: {', '.join(missing_cols)}"
        logger.error(error_msg)
        raise KeyError(error_msg)
    
    logger.debug(f"处理后的股票数据列: {list(stock_df.columns)}")
    return stock_df

# 获取所有A股股票信息
def get_all_a_stocks():
    """
    获取所有A股股票信息，并添加市值和股价数据
    """
    try:
        logger.info("开始获取所有A股股票信息...")
        start_time = time.time()
        
        # 尝试获取A股股票列表
        try:
            # 使用akshare获取A股股票列表
            stock_df = ak.stock_zh_a_spot_em()
            logger.info(f"成功获取A股股票列表，共 {len(stock_df)} 只股票")
            
            # 处理列名
            stock_df = process_stock_columns(stock_df)
            
            # 提取股票代码列表
            stock_codes = stock_df['代码'].tolist()
            
        except Exception as e:
            logger.error(f"获取A股股票列表时出错: {e}")
            logger.warning("尝试使用备用方法获取股票列表...")
            
            # 备用方法：分别获取上交所和深交所的股票列表
            try:
                # 获取上交所股票列表
                sh_stocks = ak.stock_sh_a_spot()
                logger.debug(f"获取到上交所股票 {len(sh_stocks)} 只")
                
                # 获取深交所股票列表
                sz_stocks = ak.stock_sz_a_spot()
                logger.debug(f"获取到深交所股票 {len(sz_stocks)} 只")
                
                # 处理列名
                sh_stocks = process_stock_columns(sh_stocks)
                sz_stocks = process_stock_columns(sz_stocks)
                
                # 合并数据
                stock_df = pd.concat([sh_stocks, sz_stocks], ignore_index=True)
                logger.info(f"成功合并上交所和深交所股票列表，共 {len(stock_df)} 只股票")
                
                # 提取股票代码列表
                stock_codes = stock_df['代码'].tolist()
                
            except Exception as e2:
                logger.error(f"备用方法获取股票列表也失败: {e2}")
                raise RuntimeError("无法获取A股股票列表")
        
        # 批量获取市值和股价数据
        market_data = get_batch_stock_data(stock_codes)
        
        # 合并股票信息和市值数据
        logger.info("合并股票信息和市值数据...")
        result_df = pd.merge(stock_df, market_data, on='代码', how='left')
        
        # 计算执行时间
        exec_time = time.time() - start_time
        logger.info(f"获取所有A股股票信息完成，共 {len(result_df)} 条记录，耗时 {exec_time:.2f} 秒")
        
        # 统计数据
        valid_data_count = result_df['现价'].notna().sum()
        logger.info(f"成功获取 {valid_data_count} 只股票的完整数据，{len(result_df) - valid_data_count} 只股票数据不完整")
        
        return result_df
        
    except Exception as e:
        logger.critical(f"获取所有A股股票信息时发生严重错误: {e}")
        raise

# 主程序
def main():
    """
    主程序入口
    """
    try:
        logger.info("程序开始运行...")
        
        # 获取所有A股股票信息
        stock_data = get_all_a_stocks()
        
        # 添加筛选条件：市值小于100亿且股价小于10元
        logger.info("应用筛选条件：市值小于100亿且股价小于10元...")
        
        # 确保数据类型正确
        stock_data_filtered = stock_data.copy()
        
        # 尝试进行筛选
        try:
            # 市值筛选：小于100亿（注意单位是元，需要转换为亿元）
            market_cap_filter = (stock_data_filtered['总市值_x'] / 100000000) < 100
            # 股价筛选：小于10元
            price_filter = stock_data_filtered['现价'] < 10
            # 应用组合筛选条件
            stock_data_filtered = stock_data_filtered[market_cap_filter & price_filter]
            
            logger.info(f"筛选完成，符合条件的股票数量：{len(stock_data_filtered)}")
            logger.info(f"筛选前股票数量：{len(stock_data)}，筛选后占比：{len(stock_data_filtered)/len(stock_data)*100:.2f}%")
            
            # 如果筛选后数据为空，发出警告
            if stock_data_filtered.empty:
                logger.warning("没有符合筛选条件的股票，请检查数据和筛选条件")
                stock_data_filtered = stock_data  # 如果没有符合条件的数据，返回原始数据
                logger.info("将使用原始数据进行保存")
                
        except KeyError as e:
            logger.error(f"筛选时缺少必要的列: {e}")
            stock_data_filtered = stock_data  # 如果缺少必要列，返回原始数据
        except Exception as e:
            logger.error(f"筛选过程中出错: {e}")
            stock_data_filtered = stock_data  # 如果筛选失败，返回原始数据
        
        # 创建结果目录
        result_dir = 'stock_results'
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        
        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        csv_file = os.path.join(result_dir, f'all_a_stocks_with_market_data_filtered_{timestamp}.csv')
        excel_file = os.path.join(result_dir, f'all_a_stocks_with_market_data_filtered_{timestamp}.xlsx')
        
        # 保存为CSV文件
        logger.info(f"保存筛选后数据到CSV文件: {csv_file}")
        stock_data_filtered.to_csv(csv_file, index=False, encoding='utf-8-sig')
        
        # 保存为Excel文件
        logger.info(f"保存筛选后数据到Excel文件: {excel_file}")
        stock_data_filtered.to_excel(excel_file, index=False)
        
        logger.info(f"数据保存完成，CSV文件行数: {len(stock_data_filtered)}")
        logger.info("程序运行结束")
        
    except Exception as e:
        logger.critical(f"程序运行出错: {e}")
        print(f"错误: {e}")
        return 1
    
    return 0

if __name__ == "__main__":
    exit_code = main()
    exit(exit_code)